#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>

#include "utils.h"


int add_matrix(struct matrix_t* a, struct matrix_t* b);

int main(int argc, char** argv)
{
	struct MPI_Info info;
	struct matrix_t* mat;
	struct matrix_t* smat;
	struct matrix_t* ans_mat;
	struct matrix_t* aux_mat;
	long* rows;
	long* b_col;
	long* ans;
	long i, j;
	long chunk_size, a_rows, b_rows, b_cols;

	MPI_Init(&argc, &argv);
	get_mpi_info(&info);

	if (info.proc_id == 0) {
		read_matrix("A.txt", &mat);
		read_matrix("B.txt", &smat);
		alloc_matrix(&ans_mat, mat->m, smat->n);

		chunk_size = mat->m / info.num_procs;
		chunk_size *= mat->n;
		MPI_Bcast(&chunk_size, 1, MPI_LONG, 0, MPI_COMM_WORLD);
		MPI_Bcast(&(mat->m), 1, MPI_LONG, 0, MPI_COMM_WORLD);
		MPI_Bcast(&(smat->m), 1, MPI_LONG, 0, MPI_COMM_WORLD);
		MPI_Bcast(&(smat->n), 1, MPI_LONG, 0, MPI_COMM_WORLD);
		a_rows = mat->m;
		b_rows = smat->m;
		b_cols = smat->n;
		ans = (long*) malloc(sizeof(long)*chunk_size*b_cols/b_rows);
		MPI_Scatter(
			mat->mat,
			chunk_size,
			MPI_LONG,
			MPI_IN_PLACE,
			0,
			MPI_LONG,
			0,
			MPI_COMM_WORLD);
	} else {
		MPI_Bcast(&chunk_size, 1, MPI_LONG, 0, MPI_COMM_WORLD);
		MPI_Bcast(&a_rows, 1, MPI_LONG, 0, MPI_COMM_WORLD);
		MPI_Bcast(&b_rows, 1, MPI_LONG, 0, MPI_COMM_WORLD);
		MPI_Bcast(&b_cols, 1, MPI_LONG, 0, MPI_COMM_WORLD);
		rows = (long*) malloc(sizeof(long)*chunk_size);
		b_col = (long*) malloc(sizeof(long)*b_rows);
		ans = (long*) malloc(sizeof(long)*chunk_size*b_cols/b_rows);
		MPI_Scatter(
			NULL,
			chunk_size,
			MPI_LONG,
			rows,
			chunk_size,
			MPI_LONG,
			0,
			MPI_COMM_WORLD);
	}

	if (info.proc_id == 0) {
		for (i = 0; i < b_cols; ++i) {
			sub_matrix(smat, &aux_mat, 0, smat->m-1, i, i);
			MPI_Bcast(aux_mat->mat, smat->m, MPI_LONG, 0, MPI_COMM_WORLD);
			for (j = 0; j < chunk_size/b_rows; ++j) {
				mac(mat->mat+j*b_rows, aux_mat->mat, b_rows, ans+i+j*b_cols);
			}
		}
	} else {
		for (i = 0; i < b_cols; ++i) {
			MPI_Bcast(b_col, b_rows, MPI_LONG, 0, MPI_COMM_WORLD);
			for (j = 0; j < chunk_size/b_rows; ++j) {
				mac(rows+j*b_rows, b_col, b_rows, ans+i+j*b_cols);
			}
		}
	}

	if (info.proc_id == 0) {
		free(aux_mat);
		alloc_matrix(&aux_mat, mat->m, smat->n);
	}
	MPI_Gather(
		ans, 
		chunk_size*b_cols/b_rows, 
		MPI_LONG,
		aux_mat->mat,
		chunk_size*b_cols/b_rows,
		MPI_LONG,
		0,
		MPI_COMM_WORLD);

	for (i = 0; i < aux_mat->m && info.proc_id == 0; ++i) {
		for (j = 0; j < aux_mat->n; ++j) {
			printf("%ld\t", get_elem(aux_mat, i, j));
		}
		printf("\n");
	}
			
	MPI_Barrier(MPI_COMM_WORLD);
	MPI_Finalize();

	return 0;
}
